use axum::{
    async_trait,
    extract::FromRequestParts,
    headers::{authorization::Bearer, Authorization},
    http::{request::Parts, StatusCode},
    response::IntoResponse,
    Json, RequestPartsExt, TypedHeader,
};
use jsonwebtoken as jwt;
use jwt::{DecodingKey, EncodingKey, Validation};
use once_cell::sync::Lazy;
use serde::{Deserialize, Serialize};

// secret key
static KEYS: Lazy<Keys> = Lazy::new(|| {
    // let secret = std::env::var("JWT_SECRET").expect("JWT_SECRET must be set");
    Keys::new("axum-jwt")
});

#[derive(Debug, Serialize, Deserialize)]
pub struct Claims {
    pub id: usize,
    pub username: String,
    pub exp: usize,
}

impl Claims {
    pub fn new(id: usize, username: &str) -> Self {
        Claims {
            id,
            username: username.to_string(),
            exp: 2000000000,
        }
    }

    pub fn get_jwt(&self) -> String {
        jwt::encode(&jwt::Header::default(), &self, &KEYS.encoding_key).unwrap()
    }
}

// #[async_trait]
// impl<S, B> FromRequest<S, B> for Claims
// where
//     // these bounds are required by `async_trait`
//     B: Send + 'static,
//     S: Send + Sync,
// {
//     type Rejection = HttpError;

//     async fn from_request(req: Request<B>, state: &S) -> Result<Self, Self::Rejection> {
//         let TypedHeader(Authorization(bearer)) =
//             TypedHeader::<Authorization<Bearer>>::from_request(req, state)
//                 .await
//                 .map_err(|_| HttpError::Auth)?;

//         let key = jwt::DecodingKey::from_secret(SECRET);
//         let token = jwt::decode::<Claims>(bearer.token(), &key, &Validation::default())
//             .map_err(|_| HttpError::Auth)
//             .unwrap();
//         Ok(token.claims)
//     }
// }

#[async_trait]
impl<S> FromRequestParts<S> for Claims
where
    S: Send + Sync,
{
    type Rejection = HttpError;

    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
        // Extract the token from the authorization header
        let TypedHeader(Authorization(bearer)) = parts
            .extract::<TypedHeader<Authorization<Bearer>>>()
            .await
            .map_err(|_| HttpError::Auth)?;
        let token_data =
            jwt::decode::<Claims>(bearer.token(), &KEYS.decoding_key, &Validation::default())
                .map_err(|_| HttpError::Auth)?;

        Ok(token_data.claims)
    }
}

struct Keys {
    encoding_key: EncodingKey,
    decoding_key: DecodingKey,
}

impl Keys {
    pub fn new(secret: &str) -> Self {
        let b = secret.as_bytes();
        let encoding_key = jwt::EncodingKey::from_secret(b);
        let decoding_key = jwt::DecodingKey::from_secret(b);
        Keys {
            encoding_key,
            decoding_key,
        }
    }
}

#[derive(Debug)]
pub enum HttpError {
    Auth,
    Login,
    Internal,
}

impl IntoResponse for HttpError {
    fn into_response(self) -> axum::response::Response {
        let (status, error_message) = match self {
            HttpError::Auth => (StatusCode::UNAUTHORIZED, "Unauthorized"),
            HttpError::Login => (StatusCode::FORBIDDEN, "Email or password is Error!"),
            HttpError::Internal => (StatusCode::INTERNAL_SERVER_ERROR, "Internal Server Error"),
        };
        let body = Json(serde_json::json!({
            "error": error_message,
        }));
        (status, body).into_response()
    }
}
